import cv2
import numpy as np
from TOOL import average, fenkuai


def image_concat(divide_image):
    m, n, grid_h, grid_w = [divide_image.shape[0], divide_image.shape[1],  # 每行，每列的图像块数
                            divide_image.shape[2], divide_image.shape[3]]  # 每个图像块的尺寸
    print("m,n为：" + str(grid_h) + "   " + str(grid_w))

    restore_image = np.zeros([m * grid_h, n * grid_w], np.uint8)
    # restore_image[0:grid_h, 0:]
    for i in range(m):
        for j in range(n):
            restore_image[i * grid_h:(i + 1) * grid_h, j * grid_w:(j + 1) * grid_w] = divide_image[i, j]
            plotPath = str(i) + str(j) + '.jpg'  # 图片保存路径
            # cv2.imwrite("./" + str(3) + "/" + plotPath, restore_image)

            # print(restore_image)
    return restore_image


global image_divide


def implant(image_initial, bit_stream, m=4, n=3):
    sixteen_block = fenkuai.qiege(m, n, image_initial, 1)  # 将当前块的图像分割为4*4部分
    # print(sixteen_block.shape)
    bit = str(bit_stream)

    bit_index = 0
    bit_stream=[]
    for i in range(m):
        for j in range(n):
            image_gray = sixteen_block[i, j, :]  # 分别赋予dA,dB两个变量对应图像

            da, db, dA, dB, img1, img2, bit_end = average.bit_insert(image_gray, int(bit[bit_index]))  # 小写的是旧版像素，大写的是新版像素
            bit_stream.append(bit_end)
            bit_index = bit_index + 1
            print(str(i + j) + "张  修改前的平均值da，db为：" + str(da) + "   " + str(db) + "     修改后的平均值dA，dB为" + str(
                dA) + "   " + str(dB) + str(img1
                                            .shape))

            image_newA = average.average_recover(img1, da, dA)
            image_newB = average.average_recover(img2, db, dB)
            image_new = cv2.vconcat([image_newA, image_newB])

            if i == 0 and j == 0:
                h, w = image_new.shape[0], image_new.shape[1]  # 将处理完的图片存入image_divide数组中
                image_divide = np.zeros([m, n, h, w], np.uint8)
            image_divide[i, j] = image_new
            # print(image_divide[i, j].shape)
    # print(bit_stream)
    restore_image = image_concat(image_divide)
    # cv2.imwrite("./restire.jpg", restore_image)


